{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Find the Punishment Number of an Integer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Medium"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #math #backtracking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #数学 #回溯"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: punishmentNumber"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #求一个整数的惩罚数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你一个正整数&nbsp;<code>n</code>&nbsp;，请你返回&nbsp;<code>n</code>&nbsp;的&nbsp;<strong>惩罚数</strong>&nbsp;。</p>\n",
    "\n",
    "<p><code>n</code>&nbsp;的 <strong>惩罚数</strong>&nbsp;定义为所有满足以下条件 <code>i</code>&nbsp;的数的平方和：</p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= i &lt;= n</code></li>\n",
    "\t<li><code>i * i</code> 的十进制表示的字符串可以分割成若干连续子字符串，且这些子字符串对应的整数值之和等于 <code>i</code> 。</li>\n",
    "</ul>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>n = 10\n",
    "<b>输出：</b>182\n",
    "<b>解释：</b>总共有 3 个整数 i 满足要求：\n",
    "- 1 ，因为 1 * 1 = 1\n",
    "- 9 ，因为 9 * 9 = 81 ，且 81 可以分割成 8 + 1 。\n",
    "- 10 ，因为 10 * 10 = 100 ，且 100 可以分割成 10 + 0 。\n",
    "因此，10 的惩罚数为 1 + 81 + 100 = 182\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>n = 37\n",
    "<b>输出：</b>1478\n",
    "<b>解释：</b>总共有 4 个整数 i 满足要求：\n",
    "- 1 ，因为 1 * 1 = 1\n",
    "- 9 ，因为 9 * 9 = 81 ，且 81 可以分割成 8 + 1 。\n",
    "- 10 ，因为 10 * 10 = 100 ，且 100 可以分割成 10 + 0 。\n",
    "- 36 ，因为 36 * 36 = 1296 ，且 1296 可以分割成 1 + 29 + 6 。\n",
    "因此，37 的惩罚数为 1 + 81 + 100 + 1296 = 1478\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= n &lt;= 1000</code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [find-the-punishment-number-of-an-integer](https://leetcode.cn/problems/find-the-punishment-number-of-an-integer/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [find-the-punishment-number-of-an-integer](https://leetcode.cn/problems/find-the-punishment-number-of-an-integer/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['10', '37']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        \n",
    "        ans = 0\n",
    "        \n",
    "        def check(x, cur, target):\n",
    "            \n",
    "            if cur > target:\n",
    "                return False\n",
    "            if not x:\n",
    "                return cur == target\n",
    "            \n",
    "            for i in range(1, len(x)+1):\n",
    "                if check(x[i:], cur + int(x[:i]), target):\n",
    "                    return True\n",
    "            return False\n",
    "        \n",
    "        \n",
    "        for i in range(1, n+1):\n",
    "            sq = pow(i, 2)\n",
    "            if sq == i: \n",
    "                ans += i\n",
    "                continue\n",
    "            sq = str(sq)\n",
    "            \n",
    "            if check(sq, 0, i):\n",
    "                ans += int(sq)\n",
    "        \n",
    "        return ans\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        def dfs(s,num):\n",
    "            if len(s)==0:\n",
    "                return num==0\n",
    "            for i in range(1,len(s)+1):\n",
    "                if dfs(s[i:],num-int(s[:i])):\n",
    "                    return True\n",
    "            return False\n",
    "        ans=0\n",
    "        for i in range(1,n+1):\n",
    "            num=i*i\n",
    "            if dfs(str(num),i):\n",
    "                ans+=num\n",
    "        return ans\n",
    "        \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "Pre_sum = [0] * 1001\n",
    "for i in range(1,1001):\n",
    "    s = str(i*i)\n",
    "    n = len(s)\n",
    "    def dfs(pos , sums ) -> bool:\n",
    "        if pos == n :\n",
    "            return sums == i\n",
    "        x = 0\n",
    "        for j in range(pos , n):\n",
    "            x = x*10 + int(s[j])\n",
    "            if dfs(j + 1 , sums + x):\n",
    "                return True\n",
    "        return False\n",
    "    Pre_sum[i] = Pre_sum[i - 1] + (i * i if  dfs(0,0) else 0 )\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        return Pre_sum[n]\n",
    "   \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        # def dfs(s: str, pos: int, tot: int, target: int) -> bool:\n",
    "        #     if pos == len(s):\n",
    "        #         return tot == target\n",
    "        #     sum = 0\n",
    "        #     for i in range(pos, len(s)):\n",
    "        #         sum = sum * 10 + int(s[i])\n",
    "        #         if sum + tot > target:\n",
    "        #             break\n",
    "        #         if dfs(s, i + 1, sum + tot, target):\n",
    "        #             return True\n",
    "        #     return False\n",
    "        # res = 0\n",
    "        # for i in range(1, n + 1):\n",
    "        #     if dfs(str(i * i), 0, 0, i):\n",
    "        #         res += i * i\n",
    "        # return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        ans = set()\n",
    "        \n",
    "        for number in range(1,n+1):\n",
    "            tot = 0\n",
    "            num = str(number*number)\n",
    "            m = len(num)\n",
    "            def backtrace(i):\n",
    "                nonlocal tot\n",
    "                if i == m:\n",
    "                    if tot == number:\n",
    "                        ans.add(int(num))\n",
    "                    return\n",
    "                # pick\n",
    "                for j in range(i,m):\n",
    "                    sub = int(num[i:j+1])\n",
    "                    tot += sub\n",
    "                    backtrace(j+1)\n",
    "                    tot -= sub\n",
    "                    if sub == '0' : break                \n",
    "            backtrace(0)\n",
    "        return sum(ans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        ans = set()\n",
    "        \n",
    "        for number in range(1,n+1):\n",
    "            tot = 0\n",
    "            num = str(number*number)\n",
    "            m = len(num)\n",
    "            def backtrace(i):\n",
    "                nonlocal tot\n",
    "                if tot > number: return\n",
    "                if i == m:\n",
    "                    if tot == number:\n",
    "                        ans.add(int(num))\n",
    "                    return\n",
    "                # pick\n",
    "                for j in range(i,m):\n",
    "                    sub = int(num[i:j+1])\n",
    "                    tot += sub\n",
    "                    backtrace(j+1)\n",
    "                    tot -= sub\n",
    "                    if sub == '0' : break                \n",
    "            backtrace(0)\n",
    "        return sum(ans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "\n",
    "        def check(t):\n",
    "            s=str(t*t)\n",
    "            m=len(s)\n",
    "            @cache\n",
    "            def dfs(cur,ps):\n",
    "                if cur==m:\n",
    "                    return ps==t\n",
    "                sum=0\n",
    "                for k in range(cur,m):\n",
    "                    sum=sum*10+int(s[k])\n",
    "                    if ps+sum<=t:\n",
    "                        if dfs(k+1,ps+sum):\n",
    "                            return True\n",
    "                return False\n",
    "\n",
    "            return dfs(0,0)\n",
    "\n",
    "\n",
    "        res=0\n",
    "        for i in range(1,n+1):\n",
    "            j=i*i\n",
    "            if check(i):\n",
    "                print(\"ad\",i)\n",
    "                res+=j\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "\n",
    "        def check(t):\n",
    "            s=str(t*t)\n",
    "            m=len(s)\n",
    "            @cache\n",
    "            def dfs(cur,ps):\n",
    "                if cur==m:\n",
    "                    return ps==t\n",
    "                sum=0\n",
    "                for k in range(cur,m):\n",
    "                    sum=sum*10+int(s[k])\n",
    "                    if ps+sum>t:break\n",
    "                    if dfs(k+1,ps+sum):\n",
    "                        return True\n",
    "                return False\n",
    "\n",
    "            return dfs(0,0)\n",
    "\n",
    "\n",
    "        res=0\n",
    "        for i in range(1,n+1):\n",
    "            j=i*i\n",
    "            if check(i):\n",
    "                res+=j\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        def check(num, square):\n",
    "            str_square = str(square)\n",
    "            length = len(str_square)\n",
    "            @cache\n",
    "            def helper(cnt, index):\n",
    "                if cnt == 0 and index == length:\n",
    "                    return True\n",
    "                if cnt < 0 or index == length:\n",
    "                    return False\n",
    "                for i in range(index, length):\n",
    "                    if helper(cnt - int(str_square[index: i + 1]), i+1):\n",
    "                        return True\n",
    "                return False\n",
    "            \n",
    "            return helper(num, 0)\n",
    "        \n",
    "        res = []\n",
    "        for i in range(1, n + 1):\n",
    "            if check(i, i * i):\n",
    "                res.append(i * i)\n",
    "        return sum(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def is_valid(self, i):\n",
    "        s = str(i * i)\n",
    "\n",
    "        @cache\n",
    "        def dfs(start, summ):\n",
    "            if start == len(s):\n",
    "                return summ == i\n",
    "            \n",
    "            for end in range(start, len(s)):\n",
    "                number = int(s[start: end + 1])\n",
    "                if dfs(end + 1, summ + number):\n",
    "                    return True\n",
    "            \n",
    "            return False\n",
    "\n",
    "        return dfs(0, 0)\n",
    "\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "\n",
    "        ans = 0\n",
    "        for i in range(1, n + 1):\n",
    "            if self.is_valid(i):\n",
    "                ans += i * i\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "@cache\n",
    "def isValid(acc, x):\n",
    "    if acc == x:\n",
    "        return True\n",
    "    d = 10\n",
    "    while acc // d:\n",
    "        if isValid(acc // d, x - (acc % d)):\n",
    "            return True \n",
    "        d *= 10\n",
    "    return False\n",
    "\n",
    "f = [0] * 1001\n",
    "for i in range(1001):\n",
    "    f[i] = f[i - 1] + (i * i if isValid(i * i, i) else 0)\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        return f[n]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "@cache\n",
    "def check(s, i):\n",
    "    if not s:\n",
    "        return not i\n",
    "    for j in range(1, len(str(i)) + 1):\n",
    "        if i < int(s[:j]):\n",
    "            return False\n",
    "        if check(s[j:], i - int(s[:j])):\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "punish_nums = []\n",
    "for x in range(1, 1001):\n",
    "    if check(str(x * x), x):\n",
    "        punish_nums.append(x)\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        ans = 0\n",
    "        for x in range(1, n + 1):\n",
    "            if x in punish_nums:\n",
    "                ans += x * x\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        def is_str_sum(s, target):\n",
    "            if target<0:\n",
    "                return False\n",
    "            if int(s)==target:\n",
    "                return True\n",
    "            else:\n",
    "                n = len(s)\n",
    "                for i in range(1, n):\n",
    "                    a1 = int(s[:i])\n",
    "                    res = is_str_sum(s[i:], target-a1)\n",
    "                    if res:\n",
    "                        return True\n",
    "                return False\n",
    "\n",
    "        @cache\n",
    "        def dp(n):\n",
    "            if n==1:\n",
    "                return 1\n",
    "            \n",
    "            res = dp(n-1)\n",
    "            i2 = n*n \n",
    "            if is_str_sum(str(i2), n):\n",
    "                res += i2\n",
    "            \n",
    "            return res\n",
    "        \n",
    "        return dp(n)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "arr = set()\n",
    "mx = 1001\n",
    "@cache\n",
    "def check(x, s):\n",
    "    if x == \"\" and s == 0:\n",
    "        return True\n",
    "    if x == \"\" or s < 0:\n",
    "        return False\n",
    "    m = len(x)\n",
    "    ans = False\n",
    "    for i in range(1,m+1):\n",
    "        # s += int(x[:i])\n",
    "        rest = x[i:]\n",
    "        ans |= check(rest, s - int(x[:i]))\n",
    "    return ans\n",
    "\n",
    "for i in range(1, 1001):\n",
    "    x = i * i\n",
    "    if check(str(x), i):\n",
    "        arr.add(x)\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        # print(arr)\n",
    "        ans = 0\n",
    "        for i in range(1, n + 1):\n",
    "            if i * i in arr:\n",
    "                ans += i * i\n",
    "        return ans\n",
    "                \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "@cache\n",
    "def check(s: str, target: int) -> True:   # s中是否存在某种分配方式使其和为i\n",
    "    if len(s) == 0 and target == 0:\n",
    "        return True\n",
    "    for i in range(0, len(s)):\n",
    "        if check(s[i+1:], target - int(s[:i+1])):\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        return sum([i * i for i in range(1, n + 1) if check(str(i*i), i)])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        def check(i: int) -> int:\n",
    "            pow = i**2\n",
    "            s = str(pow)\n",
    "\n",
    "            @cache\n",
    "            def dfs(l: int, r: int, target: int) -> bool:\n",
    "                if l < 0:\n",
    "                    return r == 0 and target == 0\n",
    "                # [l, r) 是否分割\n",
    "                return dfs(l - 1, l, target - int(s[l:r])) or dfs(l - 1, r, target)\n",
    "\n",
    "            n = len(s)\n",
    "            return pow if dfs(n - 1, n, i) else 0\n",
    "\n",
    "        return sum(check(i) for i in range(1, n + 1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "@cache\n",
    "def check(i):\n",
    "    s = str(i * i)\n",
    "    return dfs(s, 0, i)\n",
    "\n",
    "@cache\n",
    "def dfs(s, i, target):\n",
    "    if i == len(s) and target == 0:\n",
    "        return True\n",
    "    cur = 0\n",
    "    for idx in range(i, len(s)):\n",
    "        cur = cur * 10 + int(s[idx])\n",
    "        if cur > target:\n",
    "            return False\n",
    "        if dfs(s, idx + 1, target - cur):\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        res = 0\n",
    "\n",
    "        for i in range(1, n + 1):\n",
    "            if check(i):\n",
    "                res += i * i\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "d = {}\n",
    "\n",
    "def check(x, y, s):\n",
    "    if not y:\n",
    "        return x == s\n",
    "    \n",
    "    if (x,y,s) in d:\n",
    "        return d[(x,y,s)]\n",
    "    \n",
    "    flag = False\n",
    "    for i in range(1,len(y)+1):\n",
    "        if flag:\n",
    "            break\n",
    "        if check(x, y[i:], s+int(y[:i])):\n",
    "            flag = True\n",
    "\n",
    "    d[(x,y,s)] = flag\n",
    "    return d[(x,y,s)]\n",
    "\n",
    "pool = []\n",
    "\n",
    "for i in range(1,1001):\n",
    "    if check(i, str(i*i), 0):\n",
    "        pool.append(i)\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "\n",
    "        res = 0\n",
    "        pos = bisect.bisect_right(pool,n)\n",
    "        for i in range(pos):\n",
    "            res += pool[i] * pool[i]\n",
    "        \n",
    "        return res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "@cache\n",
    "def f(s, idx, k):\n",
    "    if not idx:\n",
    "        return k == 0\n",
    "    if k < 0:\n",
    "        return False\n",
    "    for i in range(idx-1, -1, -1):\n",
    "        if f(s, i, k - int(s[i:idx])):\n",
    "            return True\n",
    "    return False\n",
    "st = [0]\n",
    "sst = []\n",
    "for i in range(1, 1001):\n",
    "    s = str(i * i)\n",
    "    if f(s, len(s), i):\n",
    "        st.append(st[-1] + i * i)\n",
    "        sst.append(i)\n",
    "\n",
    "class Solution:\n",
    "    def punishmentNumber(self, n: int) -> int:\n",
    "        return st[bisect_right(sst, n)]"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
